﻿#ifndef __SOFTMAX_TRAINER_H__
#define __SOFTMAX_TRAINER_H__
#include "layers/Layer.h"
#include "DataSet.h"
#include "Trainer.h"

class SoftMaxTrainer : public Trainer
{
public:
    SoftMaxTrainer()
    {
        m_predict_classify_index = -1;
        m_predict_probability    = 0.f;
    }

    std::string Predict(std::shared_ptr<DataImage> input_image)
    {
        auto start_time = get_timestamp();
        if (input_image == nullptr) return "";
        auto net_image = (input_image->GetWidth() == GetNet()->GetWidth() && input_image->GetHeight() == GetNet()->GetHeight()) ? input_image : input_image->ScaleFill(GetNet()->GetWidth(), GetNet()->GetHeight());
        auto results   = Trainer::Predict(net_image->GetData(), net_image->GetWidth(), net_image->GetHeight());
        auto size      = results->size();
        auto max_index = 0;
        auto max_value = results->at(0);

        for (int i = 0; i < size; i++)
        {
            if (results->at(i) > max_value)
            {
                max_index = i;
                max_value = results->at(i);
            }
            // std::cout << i << " " << results->at(i) << std::endl;
        }
        auto end_time            = get_timestamp();
        m_predict_classify_index = max_index;
        m_predict_probability    = max_value;
        std::cout << "predict use time: " << (end_time - start_time) << "ms" << ", classify: " << GetIndexClassify(max_index) << ", probability: " << m_predict_probability << std::endl;
        return GetIndexClassify(max_index);
    }

    void LoadDataItems()
    {
        ClearClassifyIndex();

        for (auto &item : m_path_classify_map)
        {
            CreateGetClassifyIndex(item.second);
        }

        auto dataitems     = GetDataSet()->GetDataItems();
        auto classify_size = GetClassifySize();

        for (auto &item : m_path_classify_map)
        {
            std::string image_path = item.first;
            int classify_index     = CreateGetClassifyIndex(item.second);

            auto dataitem = std::make_shared<DataItem>();
            auto img      = std::make_shared<DataImage>();
            if (!img->Load(image_path))
            {
                std::cout << "invalid image path: " << image_path << std::endl;
                continue;
            }

            auto truth                            = std::make_shared<Truth>(classify_size);
            *(truth->GetDatas() + classify_index) = 1.f;

            dataitem->SetImage(img);
            dataitem->GetTruths()->push_back(truth);
            dataitem->SetLoaded(true);
            dataitems->push_back(dataitem);
        }
    }

    bool SaveModelToFile(const std::string &model_path)
    {
        std::ofstream ofs(model_path, std::ios::binary | std::ios::out);
        if (!ofs) return false;
        std::ostringstream oss;
        SaveClassifyIndex(oss);
        GetNet()->GetOptions()->SetOption("softmax_classify_index", oss.str());
        return GetNet()->Save(ofs);
    }

    bool LoadModelFromFile(const std::string &model_path)
    {
        std::ifstream ifs(model_path, std::ios::binary);
        if (!ifs) return false;
        if (!GetNet()->Load(ifs)) return false;
        std::string softmax_classify_index = GetNet()->GetOptions()->GetOption("softmax_classify_index", "");
        if (!softmax_classify_index.empty())
        {
            std::istringstream iss(softmax_classify_index);
            LoadClassifyIndex(iss);
        }

        return true;
    }

    void RefreshSoftMaxLayer()
    {
        if (GetNet()->GetBackLayer()->GetOutSize() == GetClassifySize()) return;

        auto layers = GetNet()->GetLayers();
        layers->pop_back();
        layers->pop_back();

        auto connected_layer_options = std::make_shared<Options>("connected");
        connected_layer_options->SetOption("output", GetClassifySize());
        connected_layer_options->SetOption("activation", "linear");
        auto connected_layer = GetNet()->MakeLayer(GetNet()->GetBackLayer(), connected_layer_options);
        connected_layer->SetBatch(GetBatch());
        connected_layer->Resize();
        layers->push_back(connected_layer);

        auto softmax_options = std::make_shared<Options>("softmax");
        auto softmax_layer   = GetNet()->MakeLayer(GetNet()->GetBackLayer(), softmax_options);
        softmax_layer->SetBatch(GetBatch());
        softmax_layer->Resize();
        layers->push_back(softmax_layer);
    }

public:
    const std::string GetIndexClassify(int index)
    {
        return m_index_classify_map[index];
    }

    int CreateGetClassifyIndex(const std::string &classify)
    {
        auto it = m_classify_index_map.find(classify);
        if (it != m_classify_index_map.end()) return it->second;
        auto index = (int)(m_classify_index_map.size());
        m_classify_index_map.insert_or_assign(classify, index);
        m_index_classify_map.insert_or_assign(index, classify);
        return index;
    }

    void ClearClassifyIndex()
    {
        m_classify_index_map.clear();
        m_index_classify_map.clear();
    }

protected:
    std::string LoadFile(const std::string &path)
    {
        std::ifstream ifs(path, std::ios::binary);
        if (ifs)
        {
            std::string text((std::istreambuf_iterator<char>(ifs)), (std::istreambuf_iterator<char>()));
            ifs.close();
            return text;
        }
        else
        {
            std::cout << "invalid path: " << path << std::endl;
            return "";
        }
    }

    void SaveClassifyIndex(std::ostream &os)
    {
        int size = m_classify_index_map.size();
        os.write((const char *)(&size), sizeof(size));
        for (auto it = m_classify_index_map.begin(); it != m_classify_index_map.end(); it++)
        {
            auto key     = it->first;
            auto value   = it->second;
            int key_size = key.size();
            os.write((const char *)(&key_size), sizeof(key_size));
            os.write(key.c_str(), key_size);
            os.write((const char *)(&value), sizeof(value));
        }
    }

    void LoadClassifyIndex(std::istream &is)
    {
        int size = 0;
        is.read((char *)(&size), sizeof(size));
        for (int i = 0; i < size; i++)
        {
            int key_size = 0;
            std::string key;
            is.read((char *)(&key_size), sizeof(key_size));
            key.resize(key_size);
            is.read((char *)(key.c_str()), key_size);

            int value = 0;
            is.read((char *)(&value), sizeof(value));

            m_classify_index_map[key]   = value;
            m_index_classify_map[value] = key;
        }
    }

public:
    inline int GetClassifySize() { return m_classify_index_map.size(); }
    inline int GetPredictClassifyIndex() { return m_predict_classify_index; }
    inline float GetPredictProbability() { return m_predict_probability; }
    inline std::unordered_map<std::string, std::string> *GetPathClassifyMap() { return &m_path_classify_map; }

protected:
    int m_predict_classify_index;
    float m_predict_probability;
    std::unordered_map<std::string, int> m_classify_index_map;
    std::unordered_map<int, std::string> m_index_classify_map;
    std::unordered_map<std::string, std::string> m_path_classify_map;
};

#endif